package cn.fntop.utils;

import cn.fntop.context.TenantContextHolder;
import lombok.experimental.UtilityClass;
import lombok.extern.slf4j.Slf4j;

import java.util.function.Supplier;

/**
 * 租户切换工具类
 */
@Slf4j
@UtilityClass
public class TenantBroker {

    @FunctionalInterface
    public interface RunAs<T>{
        /**
         * 执行业务逻辑
         * @param tenantId
         * @throws Exception
         */
        void run(T tenantId) throws Exception;
    }

    @FunctionalInterface
    public interface ApplyAs<T,R>{
        /**
         * 执行业务逻辑,返回一个值
         * @param tenantId
         * @throws Exception
         */
        R apply(T tenantId) throws Exception;
    }

    /**
     * 以某个租户的身份运行
     * @param tenantId
     * @param func
     */
    public void runAs(Integer tenantId, RunAs<Integer> func){
        final Integer pre = TenantContextHolder.getTenantId();
        try {
            log.trace("切换租户 tenantId:{} -> {}", pre,tenantId);
            TenantContextHolder.setTenantId(tenantId);
            func.run(tenantId);
        } catch (Exception e) {
            log.error("runAs error", e);
            throw new TenantBrokerExceptionWrapper(e.getMessage(),e);
        } finally {
            log.trace("还原租户 tenantId:{} -> {}", tenantId,pre);
            TenantContextHolder.setTenantId(pre);
        }
    }

    /**
     * 以某个租户的身份运行
     * @param tenantId
     * @param func
     * @return
     * @param <T> 返回数据类型
     */
    public <T> T applyAs(Integer tenantId,ApplyAs<Integer,T> func){
        final Integer pre = TenantContextHolder.getTenantId();
        try{
            log.trace("切换租户 tenantId:{} -> {}", pre,tenantId);
            TenantContextHolder.setTenantId(tenantId);
            return func.apply(tenantId);
        } catch (Exception e) {
            log.error("applyAs error", e);
            throw new TenantBrokerExceptionWrapper(e.getMessage(),e);
        } finally {
            log.trace("还原租户 tenantId:{} -> {}", tenantId,pre);
            TenantContextHolder.setTenantId(pre);
        }
    }

    /**
     * 以某个租户的身份运行
     * @param supplier
     * @param func
     */
    public void runAs(Supplier<Integer> supplier, RunAs<Integer> func){
        runAs(supplier.get(),func);
    }

    /**
     * 以某个租户的身份运行
     * @param supplier
     * @param func
     * @return
     * @param <T> 返回数据类型
     */
    public <T> T applyAs(Supplier<Integer> supplier,ApplyAs<Integer,T> func){
        return applyAs(supplier.get(),func);
    }

    public static class TenantBrokerExceptionWrapper extends RuntimeException{
        public TenantBrokerExceptionWrapper(String message, Throwable cause) {
            super(message, cause);
        }

        public TenantBrokerExceptionWrapper(Throwable throwable) {
            super(throwable);
        }
    }
}
